import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.autograd import Variable


class poly_model(torch.nn.Module):
    def __init__(self):
        super(poly_model, self).__init__()
        self.poly = torch.nn.Linear(3, 1)  # 输入3维输出1维,设置全连接层输入输出参数(in_features,out_features,bias=True),bias默认为True

    def forward(self, x):
        out = self.poly(x)  # 激活调用上面设置的全连接层poly,传入输入全连接层的数据，返回的是对应的out_features
        return out


def make_features(x):
    x = x.unsqueeze(1)  # unsqueeze(1) 将传入的数组x中的每个元素分解成单个元素  例如:[1,2,3]=>[1],[2],[3]  [[1],[2],[3]]=>[[[1]],[[2]],[[3]]]
    return torch.cat([x ** i for i in range(1, 4)], 1)


def f(x):
    return x.mm(w_target) + b_target[0]


def get_batch(batch_size=32):
    random = torch.randn(batch_size)
    x = make_features(random)
    y = f(x)
    return Variable(x), Variable(y)


w_target = torch.FloatTensor([0.5, 3, 2.4]).unsqueeze(1)  # 自定义权重[0.5, 3, 2.4]=>[[0.5], [3], [2.4]]
b_target = torch.FloatTensor([0.9])  # 自定义偏置值

x_sample = np.arange(-3, 3.1, 0.1)
y_sample = b_target[0] + w_target[0] * x_sample + w_target[1] * x_sample ** 2 + w_target[2] * x_sample ** 3
plt.plot(x_sample, y_sample, label='real curve')
# plt.show()  # 绘制出目标函数的样子

model = poly_model()
criterion = torch.nn.MSELoss()  # 选择均方误差衡量模型的好坏(损失函数)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)  # 选择优化器为简单的梯度下降

epoch = 0
while True:
    # 获取数据
    batch_x, batch_y = get_batch()
    # 给模型传入参数进行正向传播,获得预测值output
    output = model(batch_x)
    # 给损失函数传入预测值以及真正值获得loss
    loss = criterion(output, batch_y)
    # 获取均方误差得到的损失值
    print_loss = loss.data
    # 重置优化器中的梯度
    optimizer.zero_grad()
    # 损失值反向传播求梯度(如果没有上一句清零梯度的话梯度会累加)
    loss.backward()
    # 反向传播更新模型参数
    optimizer.step()
    epoch += 1
    if epoch % 200 == 0:
        print("Loss:{}  after {} batches".format(print_loss, epoch))
    if print_loss < 1e-3:
        break

# 点状图绘制出训练后模型的样子
test_x = torch.tensor(np.arange(-3, 3.1, 0.1), dtype=torch.float32)
test_x = Variable(make_features(test_x))
test_y = model(test_x)
plt.plot(test_x.data.numpy()[:, 0], test_y.data.numpy(), 'r*')
plt.show()
